package cn.iocoder.yudao.framework.common.util.tree;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.Assert;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author wyl
 * @date 2025/4/16
 * @des
 */
@Slf4j
public class TreeUtil {
    /**
     * 构造数据树 O(N)
     * <p>
     * 不含parentId本身，获得得是他的子集，适用于知道根节点，获得下级的tree
     *
     * @param list           原集合
     * @param mKey           父级被子集关联的字段,比如id
     * @param treeConnectKey 子集关联父级的字段，比如parentId
     * @param treeSortKey    同级菜单的排序字段，比如sort
     * @param consumer       添加子集的函数，如dto.setChild(list)
     * @param parentId       顶层节点父id的值
     * @return
     */
    public static <M> List<M> buildTree(List<M> list, Function<M, ?> mKey, Function<M, ?> treeConnectKey,
                                        Function<M, ? extends Comparable> treeSortKey, Consumers<M, M> consumer, Long parentId) {
        if (CollUtil.isEmpty(list)) {
            return Collections.emptyList();
        }
        Assert.notNull(mKey, "父级被子级关联的字段不能为空");
        Assert.notNull(treeConnectKey, "子级关联父级的字段不能为空");
        Assert.notNull(consumer, "消费函数不能为空");
        //线程安全类集合
        List<M> tree = Collections.synchronizedList(new ArrayList<>());
        //按上级id分组
        final Map<?, List<M>> collect = list.parallelStream().collect(Collectors.groupingBy(treeConnectKey));
        list.parallelStream().filter(m -> {
            final boolean b = !treeConnectKey.apply(m).equals(parentId);
            if (!b) {
                tree.add(m);
            }
            return b;
        }).forEach(
                //通过对象地址引用实现父子级关联，即使父级先添加了子级，子级在添加孙子级，父子孙三级也全部都会关联
                m -> {
                    if (treeSortKey != null) {
                        consumer.accept(m, CollUtil.sort(collect.get(mKey.apply(m)), java.util.Comparator.comparing(treeSortKey)));
                    } else {
                        consumer.accept(m, collect.get(mKey.apply(m)));
                    }
                }
        );
        if (treeSortKey != null) {
            //排序
            tree.sort(Comparator.comparing(treeSortKey));
        }
        return tree.parallelStream().peek(b -> consumer.accept(b, collect.get(mKey.apply(b)))).collect(Collectors.toList());
    }


    /**
     * 构造数据树 O(N)
     * <p>
     * 含parentId本身，适用于知道根节点，获得整个tree
     *
     * @param list
     * @param mKey
     * @param treeConnectKey
     * @param treeSortKey
     * @param consumer
     * @param parentId
     * @return
     */
    public static <M> M buildTreeContainRoot(List<M> list, Function<M, ?> mKey, Function<M, ?> treeConnectKey,
                                             Function<M, ? extends Comparable> treeSortKey, Consumers<M, M> consumer, Long parentId) {

        List<M> treeData = buildTree(list, mKey, treeConnectKey, treeSortKey, consumer, parentId);

        M rootItem = list.stream().filter(item -> mKey.apply(item).equals(parentId)).findFirst().orElse(null);
        if (rootItem == null) {
            log.error("构造tree结构失败，集合中未能找到root节点！rootId={}", parentId);
            return null;
        }
        consumer.accept(rootItem, treeData);
        return rootItem;
    }


    @FunctionalInterface
    public interface Consumers<M, N> {
        /**
         * 消费函数
         *
         * @param m
         * @param n
         */
        void accept(M m, List<N> n);
    }

    /**
     * 通过Map方式构建树
     * 适用于含搜索条件的tree,不一定一个顶级节点
     *
     * @param list
     * @param idExtractor
     * @param parentIdExtractor
     * @param childrenSetter
     * @param treeSortKey
     * @param <T>
     * @return
     */
    public static <T> List<T> buildTreeByMap(
            List<T> list,
            Function<T, ?> idExtractor,
            Function<T, ?> parentIdExtractor,
            Function<T, ? extends Comparable> treeSortKey,
            BiConsumer<T, List<T>> childrenSetter

    ) {
        if (idExtractor == null || parentIdExtractor == null || childrenSetter == null) {
            throw new IllegalArgumentException("Extractors and setter cannot be null");
        }

        // 使用Set减少contains操作的时间复杂度
        Set<Object> allIds = list.stream().map(idExtractor).collect(Collectors.toSet());
        Set<Object> topIds = list.stream()
                .filter(o -> !allIds.contains(parentIdExtractor.apply(o)))
                .map(idExtractor)
                .collect(Collectors.toSet());

        // 优化groupingBy的使用，直接为每个节点准备子节点列表
        Map<Object, List<T>> nodeMap = list.stream()
                .collect(Collectors.groupingBy(parentIdExtractor));

        // 设置子节点，这里假设setFieldValue已调整为接受BiConsumer以处理类型安全
        list.forEach(node -> {
            Object nodeId = idExtractor.apply(node);
            List<T> children = nodeMap.getOrDefault(nodeId, Collections.emptyList());
            if (treeSortKey != null) {
                children = (List<T>) children.parallelStream().sorted(Comparator.comparing(treeSortKey)).collect(Collectors.toList());
            }
            childrenSetter.accept(node, children);
        });

        // 过滤出根节点并返回
        return list.stream().filter(node -> topIds.contains(idExtractor.apply(node))).collect(Collectors.toList());
    }

}
